package com.frp.protocol;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;

/**
 * 消息编解码器
 * 处理FRP协议消息的序列化和反序列化
 */
public class MessageCodec {
    private static final Logger logger = LoggerFactory.getLogger(MessageCodec.class);
    private static final ObjectMapper objectMapper = new ObjectMapper();
    
    private static final int HEADER_LENGTH = 4;  // 消息头长度（消息长度字段）
    
    /**
     * 编码消息为字节数组
     */
    public static byte[] encode(Message message) throws IOException {
        try {
            String jsonStr = objectMapper.writeValueAsString(message);
            byte[] jsonBytes = jsonStr.getBytes(StandardCharsets.UTF_8);
            
            ByteBuffer buffer = ByteBuffer.allocate(HEADER_LENGTH + jsonBytes.length);
            buffer.putInt(jsonBytes.length);
            buffer.put(jsonBytes);
            
            return buffer.array();
        } catch (IOException e) {
            logger.error("消息编码失败", e);
            throw e;
        }
    }
    
    /**
     * 从字节缓冲区解码消息
     */
    public static Message decode(ByteBuffer buffer) throws IOException {
        try {
            if (buffer.remaining() < HEADER_LENGTH) {
                return null;  // 数据不足，等待更多数据
            }
            
            buffer.mark();
            int length = buffer.getInt();
            
            if (length <= 0 || length > 1024 * 1024) {  // 限制最大消息大小为1MB
                logger.error("无效的消息长度: {}", length);
                return null;
            }
            
            if (buffer.remaining() < length) {
                buffer.reset();  // 数据不足，重置位置等待更多数据
                return null;
            }
            
            byte[] jsonBytes = new byte[length];
            buffer.get(jsonBytes);
            
            String jsonStr = new String(jsonBytes, StandardCharsets.UTF_8);
            
            // 先解析出消息类型
            MessageTypeContainer typeContainer = objectMapper.readValue(jsonStr, MessageTypeContainer.class);
            MessageType messageType = typeContainer.getType();
            
            if (messageType == null) {
                logger.error("无法确定消息类型");
                return null;
            }
            
            // 根据消息类型解析具体消息
            Class<? extends Message> messageClass = getMessageClass(messageType);
            if (messageClass == null) {
                logger.error("未知的消息类型: {}", messageType);
                return null;
            }
            
            return objectMapper.readValue(jsonStr, messageClass);
        } catch (IOException e) {
            logger.error("消息解码失败", e);
            throw e;
        }
    }
    
    /**
     * 根据消息类型获取对应的类
     */
    private static Class<? extends Message> getMessageClass(MessageType type) {
        switch (type) {
            case Login:
                return LoginMessage.class;
            case LoginResp:
                return LoginRespMessage.class;
            case Heartbeat:
                return Message.class;  // 心跳包使用基础消息类
            case HeartbeatResp:
                return Message.class;
            case NewProxy:
                return NewProxyMessage.class;
            case NewProxyResp:
                return NewProxyRespMessage.class;
            case CloseProxy:
                return CloseProxyMessage.class;
            case NewWorkConn:
                return NewWorkConnMessage.class;
            case NewWorkConnResp:
                return NewWorkConnRespMessage.class;
            case Data:
                return DataMessage.class;
            case Error:
                return ErrorMessage.class;
            default:
                return null;
        }
    }
    
    /**
     * 用于临时解析消息类型的容器类
     */
    private static class MessageTypeContainer {
        private MessageType type;
        
        public MessageType getType() {
            return type;
        }
        
        public void setType(MessageType type) {
            this.type = type;
        }
    }
}